from finder import Finder
import fitz  # PyMuPD
import tools
import pytesseract
from PIL import Image
import io


class PDFFinder(Finder):

    @staticmethod
    def file_suffixes():
        return [".pdf"]

    def search_file(self, file_path, search_content):
        if Finder.should_file_path_ignore(file_path):
            return

        doc = fitz.open(file_path)
        for page_num, page in enumerate(doc):
            text = page.get_text("text")
            if not text:
                image_list = page.get_images(full=True)
                for image_index, img in enumerate(image_list):
                    # 获取图片的 XREF
                    xref = img[0]
                    # 提取图像字节
                    base_image = doc.extract_image(xref)
                    image_bytes = base_image["image"]
                    # 读取图像
                    image = Image.open(io.BytesIO(image_bytes))
                    # 使用 PyTesseract 识别图像中的文字
                    # TODO: 这里加try catch
                    try:
                        text = text + pytesseract.image_to_string(image, lang='chi_sim')
                    except Exception as e:
                        print(f"exception in image to string:{e}")
                        continue
            with tools.lock:
                tools.global_store[f'{file_path}_{page_num + 1}'] = text
            pass

            if search_content.lower() in Finder.process_text_if_need(text.lower()):
                print(f"Found '{search_content}' in {file_path} on page {page_num + 1}")
        doc.close()
